package org.example.zuul.http;

import java.util.concurrent.ConcurrentHashMap;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;

import org.springframework.http.RequestEntity;
import org.springframework.http.ResponseEntity;

public class RequestContext extends ConcurrentHashMap<String, Object> {
    private static final String KEY_REQUEST = "_request_";
    private static final String KEY_RESPONSE = "_response_";
    private static final String KEY_REQUEST_ENTITY = "_request_entity_";
    private static final String KEY_RESPONSE_ENTITY = "_response_entity_";
    protected static final ThreadLocal<? extends RequestContext> THREAD_LOCAL = ThreadLocal.withInitial(RequestContext::new);

    public HttpServletRequest getRequest() {
        return (HttpServletRequest) get(KEY_REQUEST);
    }

    public void setRequest(HttpServletRequest request) {
        put(KEY_REQUEST, request);
    }

    public HttpServletResponse getResponse() {
        return (HttpServletResponse) get(KEY_RESPONSE);
    }

    public void setResponse(HttpServletResponse response) {
        put(KEY_RESPONSE, response);
    }

    @SuppressWarnings("unchecked")
    public <T> RequestEntity<T> getRequestEntity() {
        return (RequestEntity<T>) get(KEY_REQUEST_ENTITY);
    }

    public <T> void setRequestEntity(RequestEntity<T> requestEntity) {
        put(KEY_REQUEST_ENTITY, requestEntity);
    }

    @SuppressWarnings("unchecked")
    public <T> ResponseEntity<T> getResponseEntity() {
        return (ResponseEntity<T>) get(KEY_RESPONSE_ENTITY);
    }

    public <T> void setResponseEntity(ResponseEntity<T> responseEntity) {
        put(KEY_RESPONSE_ENTITY, responseEntity);
    }

    public static RequestContext getCurrentContext() {
        return THREAD_LOCAL.get();
    }

    public void set(String key, Object value) {
        if (value != null) {
            put(key, value);
        } else {
            remove(key);
        }
    }

    public void unset() {
        THREAD_LOCAL.remove();
    }
}
